module latlon_mesh_mod

  use flogger
  use const_mod
  use sphere_geometry_mod
  use kdtree

  implicit none

  private

  public latlon_mesh_type

  type latlon_mesh_type
    type(latlon_mesh_type), pointer :: global_mesh => null()
    integer lon_halo_width
    integer lat_halo_width
    integer num_lon
    integer num_lat
    integer num_lev
    integer lon_ibeg
    integer lon_iend
    integer lat_ibeg
    integer lat_iend
    integer lat_ibeg_no_pole
    integer lat_iend_no_pole
    integer lev_ibeg
    integer lev_iend
    integer lon_lb
    integer lon_ub
    integer lat_lb
    integer lat_ub
    integer lev_lb
    integer lev_ub
    real(r8) start_lon
    real(r8) end_lon
    real(r8) start_lat
    real(r8) end_lat
    real(r8) dlon
    real(r8), allocatable, dimension(:) :: dlat
    real(r8) total_area
    real(r8), allocatable, dimension(:) :: lon
    real(r8), allocatable, dimension(:) :: lat
    real(r8), allocatable, dimension(:) :: lev
    real(r8), allocatable, dimension(:) :: cos_lon
    real(r8), allocatable, dimension(:) :: sin_lon
    real(r8), allocatable, dimension(:) :: cos_lat
    real(r8), allocatable, dimension(:) :: sin_lat
    ! For output
    real(r8), allocatable, dimension(:) :: lon_deg
    real(r8), allocatable, dimension(:) :: lat_deg
    ! Area for weighting
    real(r8), allocatable, dimension(:) :: area_cell
    ! Edge length
    real(r8), allocatable, dimension(:) :: dx
    real(r8), allocatable, dimension(:) :: dy
  contains
    procedure :: init_global => latlon_mesh_init_global
    procedure :: init_from_parent => latlon_mesh_init_from_parent
    procedure :: common_init => latlon_mesh_common_init
    procedure :: has_south_pole => latlon_mesh_has_south_pole
    procedure :: has_north_pole => latlon_mesh_has_north_pole
    procedure :: is_south_pole => latlon_mesh_is_south_pole
    procedure :: is_north_pole => latlon_mesh_is_north_pole
    procedure :: is_pole => latlon_mesh_is_pole
    procedure :: is_inside_with_halo_lat => latlon_mesh_is_inside_with_halo_lat
    procedure :: is_outside_pole_lat => latlon_mesh_is_outside_pole_lat
    final :: latlon_mesh_final
  end type latlon_mesh_type

contains

  subroutine latlon_mesh_init_global(this, lon, lat, lev)

    class(latlon_mesh_type), intent(inout), target :: this
    real(r8), intent(in) :: lon(:)
    real(r8), intent(in) :: lat(:)
    real(r8), intent(in) :: lev(:)

    integer i, j, k
    logical decrease

    this%global_mesh => this

    this%num_lon  = size(lon)
    this%lon_ibeg = 1
    this%lon_iend = this%num_lon
    this%num_lat  = size(lat)
    this%lat_ibeg = 1
    this%lat_iend = this%num_lat

    this%lon_halo_width = 1
    this%lat_halo_width = 1
    this%start_lon      = 0.0_r8
    this%end_lon        =  pi2
    this%start_lat      = -pi05
    this%end_lat        =  pi05

    this%num_lev  = size(lev)
    this%lev_ibeg = 1
    this%lev_iend = size(lev)
    this%lev_lb   = 1
    this%lev_ub   = size(lev)

    call this%common_init()

    this%lon(1:size(lon)) = lon
    this%lon_deg(1:size(lon)) = this%lon(1:size(lon)) * deg
    ! Assume zonal grid intervals are equal.
    this%dlon    = (this%end_lon - this%start_lon) / this%num_lon

    this%lat(1:size(lat)) = lat
    this%lat_deg(1:size(lat)) = this%lat(1:size(lat)) * deg
    do j = this%lat_ibeg, this%lat_iend - 1
      this%dlat(j) = this%lat(j+1) - this%lat(j)
    end do

    do i = this%lon_lb, this%lon_ub
      this%cos_lon(i) = cos(this%lon(i))
      this%sin_lon(i) = sin(this%lon(i))
    end do

    do j = this%lat_lb, this%lat_ub
      if (this%lat(j) >= -pi05 .and. this%lat(j) <= pi05) then
        this%cos_lat(j) = cos(this%lat(j))
        this%sin_lat(j) = sin(this%lat(j))
      end if
    end do

    ! Ensure the values of cos_lat and sin_lat are expected at the Poles.
    this%cos_lat(this%lat_ibeg) =  0.0_r8
    this%sin_lat(this%lat_ibeg) = -1.0_r8
    this%cos_lat(this%lat_iend) =  0.0_r8
    this%sin_lat(this%lat_iend) =  1.0_r8

    do j = this%lat_ibeg, this%lat_iend
      if (this%is_south_pole(j)) then
        this%area_cell(j) = radius**2 * this%dlon * (this%sin_lat(j) + 1.0d0)
      else if (this%is_north_pole(j)) then
        this%area_cell(j) = radius**2 * this%dlon * (1.0 - this%sin_lat(j-1))
      else
        this%area_cell(j) = radius**2 * this%dlon * (this%sin_lat(j) - this%sin_lat(j-1))
      end if
    end do

    do j = this%lat_ibeg, this%lat_iend
      this%dx(j) = radius * this%cos_lat(j) * this%dlon
    end do
    do j = this%lat_ibeg, this%lat_iend - 1
      this%dy(j) = radius * this%dlat(j)
    end do

    ! Check if levels are monotonic.
    decrease = lev(1) > lev(2)
    do k = 1, size(lev) - 1
      if ((decrease .and. lev(k) < lev(k+1)) .or. (.not. decrease .and. lev(k) > lev(k+1))) then
        call log_error('Model levels are not monotonic!')
      end if
    end do
    this%lev = lev

  end subroutine latlon_mesh_init_global

  subroutine latlon_mesh_init_from_parent(this, parent, lon_halo_width, lat_halo_width, &
                                          lon_ibeg, lon_iend, lat_ibeg, lat_iend)

    class(latlon_mesh_type), intent(inout) :: this
    class(latlon_mesh_type), intent(in), target :: parent
    integer, intent(in) :: lon_halo_width
    integer, intent(in) :: lat_halo_width
    integer, intent(in) :: lon_ibeg
    integer, intent(in) :: lon_iend
    integer, intent(in) :: lat_ibeg
    integer, intent(in) :: lat_iend

    integer i, j, k

    this%global_mesh => parent

    this%num_lon  = lon_iend - lon_ibeg + 1
    this%lon_ibeg = lon_ibeg
    this%lon_iend = lon_iend
    this%num_lat  = lat_iend - lat_ibeg + 1
    this%lat_ibeg = lat_ibeg
    this%lat_iend = lat_iend
    this%num_lev  = parent%num_lev
    this%lev_ibeg = parent%lev_ibeg
    this%lev_iend = parent%lev_iend

    this%lon_halo_width = lon_halo_width
    this%lat_halo_width = lat_halo_width
    this%start_lon      = parent%lon(lon_ibeg)
    this%end_lon        = parent%lon(lon_iend) + parent%dlon
    this%start_lat      = merge(parent%lat(lat_ibeg-1), -pi05, .not. this%has_south_pole())
    this%end_lat        = merge(parent%lat(lat_iend  ),  pi05, .not. this%has_north_pole())

    call this%common_init()

    this%dlon = parent%dlon
    do i = this%lon_lb, this%lon_ub
      this%lon(i) = parent%lon(i)
      this%lon_deg(i) = parent%lon_deg(i)
      this%sin_lon(i) = parent%sin_lon(i)
      this%cos_lon(i) = parent%cos_lon(i)
    end do

    this%dlat = parent%dlat(lbound(this%dlat, 1):ubound(this%dlat, 1))
    do j = this%lat_lb, this%lat_ub
      this%lat(j) = parent%lat(j)
      this%lat_deg(j) = parent%lat_deg(j)
      this%sin_lat(j) = parent%sin_lat(j)
      this%cos_lat(j) = parent%cos_lat(j)
      this%area_cell(j) = parent%area_cell(j)
      this%dx(j) = parent%dx(j)
      if (.not. this%is_north_pole(j)) this%dy(j) = parent%dy(j)
    end do

    this%lev = parent%lev

  end subroutine latlon_mesh_init_from_parent

  subroutine latlon_mesh_common_init(this)

    class(latlon_mesh_type), intent(inout) :: this

    this%total_area = radius**2 * (this%end_lon - this%start_lon) * (sin(this%end_lat) - sin(this%start_lat))

    this%lat_ibeg_no_pole = merge(this%lat_ibeg + 1, this%lat_ibeg, this%has_south_pole())
    this%lat_iend_no_pole = merge(this%lat_iend - 1, this%lat_iend, this%has_north_pole())

    this%lon_lb = this%lon_ibeg - this%lon_halo_width
    this%lon_ub = this%lon_iend + this%lon_halo_width
    this%lat_lb = this%lat_ibeg - this%lat_halo_width
    this%lat_ub = this%lat_iend + this%lat_halo_width
    this%lev_lb = this%lev_ibeg
    this%lev_ub = this%lev_iend

    allocate(this%dlat     (this%lat_lb:this%lat_ub-1)); this%dlat      = 0.0_r8
    allocate(this%lon      (this%lon_lb:this%lon_ub  )); this%lon       = inf
    allocate(this%lat      (this%lat_lb:this%lat_ub  )); this%lat       = inf
    allocate(this%lev      (this%lev_lb:this%lev_ub  )); this%lev       = inf
    allocate(this%cos_lon  (this%lon_lb:this%lon_ub  )); this%cos_lon   = inf
    allocate(this%sin_lon  (this%lon_lb:this%lon_ub  )); this%sin_lon   = inf
    allocate(this%cos_lat  (this%lat_lb:this%lat_ub  )); this%cos_lat   = inf
    allocate(this%sin_lat  (this%lat_lb:this%lat_ub  )); this%sin_lat   = inf
    allocate(this%lon_deg  (this%lon_lb:this%lon_ub  )); this%lon_deg   = inf
    allocate(this%lat_deg  (this%lat_lb:this%lat_ub  )); this%lat_deg   = inf
    allocate(this%area_cell(this%lat_lb:this%lat_ub  )); this%area_cell = 0.0_r8
    allocate(this%dx       (this%lat_lb:this%lat_ub  )); this%dx        = 0.0_r8
    allocate(this%dy       (this%lat_lb:this%lat_ub  )); this%dy        = 0.0_r8

  end subroutine latlon_mesh_common_init

  logical function latlon_mesh_has_south_pole(this) result(res)

    class(latlon_mesh_type), intent(in) :: this

    res = this%lat_ibeg == 1

  end function latlon_mesh_has_south_pole

  logical function latlon_mesh_has_north_pole(this) result(res)

    class(latlon_mesh_type), intent(in) :: this

    res = this%lat_iend == this%global_mesh%num_lat

  end function latlon_mesh_has_north_pole

  logical function latlon_mesh_is_south_pole(this, j) result(res)

    class(latlon_mesh_type), intent(in) :: this
    integer, intent(in) :: j

    res = this%has_south_pole() .and. j == 1

  end function latlon_mesh_is_south_pole

  logical function latlon_mesh_is_north_pole(this, j) result(res)

    class(latlon_mesh_type), intent(in) :: this
    integer, intent(in) :: j

    res = this%has_north_pole() .and. j == this%global_mesh%num_lat

  end function latlon_mesh_is_north_pole

  logical function latlon_mesh_is_pole(this, j) result(res)

    class(latlon_mesh_type), intent(in) :: this
    integer, intent(in) :: j

    res = this%is_south_pole(j) .or. this%is_north_pole(j)

  end function latlon_mesh_is_pole

  logical function latlon_mesh_is_inside_with_halo_lat(this, j) result(res)

    class(latlon_mesh_type), intent(in) :: this
    integer, intent(in) :: j

    res = j >= this%lat_lb .and. j <= this%lat_ub

  end function latlon_mesh_is_inside_with_halo_lat

  logical function latlon_mesh_is_outside_pole_lat(this, j) result(res)

    class(latlon_mesh_type), intent(in) :: this
    integer, intent(in) :: j

    res = j < 1 .or. j > this%global_mesh%num_lat

  end function latlon_mesh_is_outside_pole_lat

  subroutine latlon_mesh_final(this)

    type(latlon_mesh_type), intent(inout) :: this

    if (allocated(this%lon      )) deallocate(this%lon      )
    if (allocated(this%lat      )) deallocate(this%lat      )
    if (allocated(this%lev      )) deallocate(this%lev      )
    if (allocated(this%cos_lon  )) deallocate(this%cos_lon  )
    if (allocated(this%sin_lon  )) deallocate(this%sin_lon  )
    if (allocated(this%cos_lat  )) deallocate(this%cos_lat  )
    if (allocated(this%sin_lat  )) deallocate(this%sin_lat  )
    if (allocated(this%lon_deg  )) deallocate(this%lon_deg  )
    if (allocated(this%lat_deg  )) deallocate(this%lat_deg  )
    if (allocated(this%area_cell)) deallocate(this%area_cell)
    if (allocated(this%dx       )) deallocate(this%dx       )
    if (allocated(this%dy       )) deallocate(this%dy       )

  end subroutine latlon_mesh_final

end module latlon_mesh_mod
